#include "heap.h"
#include "status.h"
#include "memory/memory.h"

static int heap_validate_table(void* ptr, void* end, struct heap_table* table)
{
    int res = 0; 
    const size_t table_size = (size_t)(end - ptr); 
    const size_t total_blocks = table_size / TESTOS_HEAP_BLOCK_SIZE; 
    if (table->entries_amount != total_blocks)
    {
        res = -TESTOS_STATUS_INVALID_ARG; 
        goto out; 
    }

out:
    return res; 
}

static int heap_validate_alignment(void* ptr) 
{
    return ((unsigned int)ptr % TESTOS_HEAP_BLOCK_SIZE) == 0; 
}

// Function that initializes a heap
// returns an error code
// 
// heap  - a pointer to the heap
// ptr   - a pointer to the data pool
// end   - a pointer to the end of a heap
// table - a pointer to the heap table
int heap_create(struct heap* heap, void* ptr, void* end, struct heap_table* table) 
{
    int res = 0; 
    size_t table_size; 

    if (!heap_validate_alignment(ptr) || !heap_validate_alignment(end)) {
        res = -TESTOS_STATUS_INVALID_ARG; 
        goto out;
    }

    memset(heap, 0, sizeof(struct heap)); 
    heap->saddr = ptr; 
    heap->table = table; 

    res = heap_validate_table(ptr, end, table); 
    if (res < 0) {
        goto out; 
    }

    table_size = sizeof(HEAP_BLOCK_TABLE_ENTRY) * table->entries_amount; 
    memset(table->entries, HEAP_BLOCK_TABLE_ENTRY_FREE, table_size); 

out:
    return res; 
}

static const uint32_t heap_align_value_to_upper(const uint32_t val)
{
    const uint32_t mod = val % TESTOS_HEAP_BLOCK_SIZE; 
    if (mod == 0) {
        return val; 
    }

    // additional block required
    return val - mod + TESTOS_HEAP_BLOCK_SIZE;    
}

static int heap_get_entry_type(HEAP_BLOCK_TABLE_ENTRY entry) 
{
    return entry & 0x0F; 
}

const int heap_get_start_block(struct heap* heap, const uint32_t total_blocks)
{
    struct heap_table* table = heap->table; 
    int bc = 0; 
    int bs = -1; 
    size_t i; 

    for (i = 0; i < table->entries_amount; i++) {
        if (heap_get_entry_type(table->entries[i]) != HEAP_BLOCK_TABLE_ENTRY_FREE) {
            bc = 0; 
            bs = -1; 
            continue; 
        }

        // Now we know that block is free

        if (bs == -1) { // if this is a first block
            bs = i; 
        }

        bc++; 

        if (bc == total_blocks) {
            break; 
        }
    }

    if (bs == -1 ) { //|| bc != total_blocks) {
        return -TESTOS_STATUS_NOMEMORY; 
    }

    return bs; 
}

void* heap_block_to_address(struct heap* heap, const int start_block) 
{
    return heap->saddr + (start_block * TESTOS_HEAP_BLOCK_SIZE); 
}

void heap_mark_blocks_taken(struct heap* heap, const int start_block, const int total_blocks)
{
    const int end_block = (start_block + total_blocks) - 1; 
    HEAP_BLOCK_TABLE_ENTRY entry = HEAP_BLOCK_TABLE_ENTRY_TAKEN | HEAP_BLOCK_IS_FIRST; 
    if (total_blocks > 1) {
        entry |= HEAP_BLOCK_HAS_NEXT; 
    }

    for (int i = start_block; i <= end_block; i++) {
        heap->table->entries[i] = entry; 
        entry = HEAP_BLOCK_TABLE_ENTRY_TAKEN; 
        if (i != end_block - 1) {
            entry |= HEAP_BLOCK_HAS_NEXT; 
        }
    }
}

void* heap_malloc_blocks(struct heap* heap, const uint32_t total_blocks)
{
    void* address = 0; 
    const int start_block = heap_get_start_block(heap, total_blocks); 
    if (start_block < 0) {
        goto out; 
    }

    address = heap_block_to_address(heap, start_block); 
    heap_mark_blocks_taken(heap, start_block, total_blocks); 

out:
    return address; 
}



void* heap_malloc(struct heap* heap, const size_t size)
{
    size_t aligned_size; 
    uint32_t total_blocks; 

    aligned_size = heap_align_value_to_upper(size); 
    total_blocks = aligned_size / TESTOS_HEAP_BLOCK_SIZE; 

    return heap_malloc_blocks(heap, total_blocks); 
}


void heap_mark_blocks_free(struct heap* heap, const int start_block) 
{
    struct heap_table* table = heap->table; 
    const int total = (int)table->entries_amount; 

    for (int i = start_block; i < total; i++) {
        HEAP_BLOCK_TABLE_ENTRY entry = table->entries[i]; 

        table->entries[i] = HEAP_BLOCK_TABLE_ENTRY_FREE; 
        if (!(entry & HEAP_BLOCK_HAS_NEXT)) {
            break; 
        }
    }
}

int heap_address_to_block(struct heap* heap, void* ptr) 
{
    return ((int)(ptr - heap->saddr)) / TESTOS_HEAP_BLOCK_SIZE; 
}

void heap_free(struct heap* heap, void* ptr)
{
    heap_mark_blocks_free(heap, heap_address_to_block(heap, ptr)); 
}